import torch
import numpy as np
from torch import Tensor
from . import Colvar
from ..function import get_tensor, GLOBAL_DEVICE_SPONGE

class XRD3D(Colvar):
    def __init__(self, 
                 theta:float=0,
                 lamb:float=1,
                 pbc_box=None,
                 index=None,
                 qi=None,
                 s:float=1,) -> None:
        super().__init__(name='xrd3d',shape=(1,))
        self.device = GLOBAL_DEVICE_SPONGE()
        self.theta = np.array(theta).reshape(-1,1)
        self.lamb = np.array(lamb).reshape(-1,1)
        self.att_coeff = np.array(s)
        self.k0 = 4*np.pi/self.lamb*np.sin(self.theta*np.pi/180)
        self.indices = get_tensor(index, torch.int32, self.device).reshape(-1)
        self.xatom_numbers = self.indices.shape[0]
        self.qi = get_tensor(qi,torch.float32,self.device).reshape(qi.shape[0],qi.shape[1],1)
        self.pbc_box = pbc_box.to(self.device)
        nfft = pbc_box * 10 // 4 * 4
        self.set_nfft(nfft)
        self.fft_fijk = self.build_fft_fijk_list(self.nfft.detach().cpu().numpy(), self.pbc_box.detach().cpu().numpy()*10)

        ma = [1.0 / 6.0, -0.5, 0.5, -1.0 / 6.0]
        ma = torch.tensor([[ma[i], ma[j], ma[k]]for i in range(4)
                     for j in range(4) for k in range(4)], dtype=torch.float32, device=self.device)
        self.ma = ma.reshape(1, 1, 64, 3)
        mb = [0, 0.5, -1, 0.5]
        mb = torch.tensor([[mb[i], mb[j], mb[k]]for i in range(4)
                     for j in range(4) for k in range(4)], dtype=torch.float32, device=self.device)
        self.mb = mb.reshape(1, 1, 64, 3)
        mc = [0, 0.5, 0, -0.5]
        mc = torch.tensor([[mc[i], mc[j], mc[k]]for i in range(4)
                     for j in range(4) for k in range(4)], dtype=torch.float32, device=self.device)
        self.mc = mc.reshape(1, 1, 64, 3)
        md = [0, 1.0 / 6.0, 4.0 / 6.0, 1.0 / 6.0]
        md = torch.tensor([[md[i], md[j], md[k]]for i in range(4)
                     for j in range(4) for k in range(4)], dtype=torch.float32, device=self.device)
        self.md = md.reshape(1, 1, 64, 3)
        self.base_grid = torch.tensor([[i, j, k] for i in range(4) for j in range(4) for k in range(4)], dtype=torch.int32, device=self.device).reshape(1, 1, 64, 3)
        
    def set_nfft(self, nfft: Tensor):
        """set nfft"""
        self.nfft = torch.tensor(nfft, dtype=torch.int32, device=self.device).reshape((-1, 1, 3))
        self.fftx = int(torch.max(self.nfft[:,0,0]))
        self.ffty = int(torch.max(self.nfft[:,0,1]))
        self.fftz = int(torch.max(self.nfft[:,0,2]))
        if self.fftx % 4 != 0 or self.ffty % 4 != 0 or self.fftz % 4 != 0:
            raise ValueError("The FFT grid number for PME must be a multiple of 4")
        self.nfft[:,0,0] = torch.where(self.nfft[:,0,0]!=self.fftx,self.fftx,self.nfft[:,0,0])
        self.nfft[:,0,1] = torch.where(self.nfft[:,0,1]!=self.ffty,self.ffty,self.nfft[:,0,1])
        self.nfft[:,0,2] = torch.where(self.nfft[:,0,2]!=self.fftz,self.fftz,self.nfft[:,0,2])

    def build_fft_fijk_list(self, grid_dimension, pbc_box):
        volume = np.prod(pbc_box,axis=1,keepdims=True)

        # 构建B样条线插值的卷积格点
        B_spline_grid = np.zeros(*grid_dimension[0], dtype=np.float32)[None].repeat(grid_dimension.shape[0],0)
        temp_b_spline = np.array([1. / 6., 2. / 3., 1. / 6.])
        grid_serial = np.array([[i, j, k] for i in range(-1,2) for j in range(-1,2) for k in range(-1,2)]) % grid_dimension
        weights = temp_b_spline.reshape(-1,1,1)*temp_b_spline.reshape(1,-1,1)*temp_b_spline.reshape(1,1,-1)
        for i in range(B_spline_grid.shape[0]):
            B_spline_grid[i,grid_serial[i,:,0], grid_serial[i,:,1], grid_serial[i,:,2]] = weights.flatten()
        fft_B_spline=np.fft.rfftn(B_spline_grid,norm='backward',axes=(1,2,3)).reshape(B_spline_grid.shape[0],-1)
        scalor = np.pi*np.sqrt(np.pi) / volume / self.k0 / np.sqrt(self.att_coeff)/self.xatom_numbers

        grid_dimension_lower_half_plus_one = (grid_dimension[0,0] // 2) + 1
        grid_dimension_upper_half = (grid_dimension[0,0] + 1) // 2
        NzNy_Nx21_numbers = grid_dimension[0,0,0] * grid_dimension[0,0,1] * grid_dimension_lower_half_plus_one[2]

        grid_i = np.arange(NzNy_Nx21_numbers).reshape(-1,1).repeat(3,1)
        fft_fijk = np.zeros(NzNy_Nx21_numbers)
        grid_layer_numbers = grid_dimension[0,0,1]*grid_dimension_lower_half_plus_one[2]
        grid_i[:,2] = grid_i[:,2] % grid_dimension_lower_half_plus_one[2]
        grid_i[:,1] = np.int32((grid_i[:,1] % grid_layer_numbers) / grid_dimension_lower_half_plus_one[2])
        grid_i[:,0] = np.int32(grid_i[:,0] / grid_layer_numbers)
        temp = grid_i - grid_dimension_upper_half
        temp = np.right_shift(temp, 31)
        grid_fft_serial = (temp & grid_i) + ((~temp) & (grid_dimension[0] - grid_i))
        pi2ijk = 2. * np.pi * np.linalg.norm(grid_fft_serial / pbc_box[:,None], axis=-1)
        pi2ijk_2 = pi2ijk**2 + self.k0**2
        rk0 = 2.*pi2ijk*self.k0
        fft_fijk = scalor * 1. / pi2ijk*(np.exp(-0.25/self.att_coeff*(pi2ijk_2 - rk0)) - np.exp(-0.25/self.att_coeff*(pi2ijk_2 + rk0))) / (np.real(fft_B_spline))**2
        fft_fijk[:,0] = 0
        return torch.tensor(fft_fijk.reshape(-1, grid_dimension[0,0,0], grid_dimension[0,0,1], grid_dimension_lower_half_plus_one[2]),device=self.device)

    def forward(self, coordinate:Tensor, pbc_box:Tensor):

        c = coordinate[:,self.indices]
        pbc_box = self.pbc_box.reshape((-1, 1, 3))
        frac = c / pbc_box.detach() % 1.0 * self.nfft
        grid = frac.to(torch.int32)
        frac = frac - torch.floor(frac)

        # (B,A,64,3) <- (B,A,1,3) + (1,1,64,3)
        neibor_grids = grid.unsqueeze(2) - self.base_grid
        neibor_grids = neibor_grids % self.nfft.unsqueeze(2)
        neibor_grids = neibor_grids.to(self.device)

        # (B,A,64,3) <- (B,A,1,3) * (1,1,64,3)
        frac = frac.unsqueeze(2)
        neibor_q = frac * frac * frac * self.ma + frac * frac * self.mb + frac * self.mc + self.md

        # (B,A,64) <- (B,A,1) * reduce (B,A,64,3)
        neibor_q = self.qi * torch.prod(neibor_q, -1)

        # (B,A,64,4) <- concat (B,A,64,1) (B,A,64,3)
        batch_constant = torch.ones((c.shape[0], c.shape[1], 64, 1), dtype=torch.int32, device=self.device)
        batch_constant = batch_constant * torch.arange(0, c.shape[0], device=self.device).reshape(-1, 1, 1, 1)
        neibor_grids = torch.cat((batch_constant, neibor_grids), -1)

        # (B, fftx, ffty, fftz)
        q_matrix = torch.zeros([c.shape[0], self.fftx, self.ffty, self.fftz], dtype=torch.float32, device=self.device)
        neibor_grids_new = (neibor_grids*torch.tensor([self.fftx*self.ffty*self.fftz,self.ffty*self.fftz,self.fftz,1],dtype=torch.int32,device=self.device)).sum(-1).to(torch.int32)
        q_matrix = torch.index_add(q_matrix.flatten(),0,neibor_grids_new.flatten(),neibor_q.flatten()).reshape((c.shape[0], self.fftx, self.ffty, self.fftz))
        
        # # XRD 
        fq = torch.fft.rfftn(q_matrix,norm='backward',dim=(1,2,3))
        fq_fijk = fq * self.fft_fijk
        phi = torch.fft.irfftn(fq_fijk,norm='forward',dim=(1,2,3))
        intensity = phi * q_matrix
        return intensity.sum(dim=(1,2,3)).to(torch.float32).unsqueeze(-1)